from typing import AsyncGenerator, Union, Any, Tuple
from datetime import datetime, timedelta
from fastapi import Depends, HTTPException, status, Request, WebSocket
from fastapi.security import OAuth2PasswordBearer
from config import settings
from jose import jwt
from pydantic import ValidationError
from models.core import TokenPayload, SystemUser, TokenSchema, RefreshToken
from services.storage import AioRedisCache
from services.db import RedisUserSchema
import uuid

# OAuth2PasswordBearer.__call__() missing 1 required positional argument: 'request'
# @see: https://github.com/fastapi/fastapi/issues/2031
class ComposeOAuth2PasswordBearer(OAuth2PasswordBearer):
    async def __call__(self, request: Request = None, websocket: WebSocket = None):
        return await super().__call__(request or websocket)

async def deps_user_db()->AsyncGenerator[RedisUserSchema, None]:
    redis_config = {'port': settings.REDIS_PORT, 'db': settings.REDIS_DB}
    if settings.REDIS_USERNAME and settings.REDIS_PASSWORD:
        redis_config['username'] = settings.REDIS_USERNAME
        redis_config['password'] = settings.REDIS_PASSWORD

    async with AioRedisCache(settings.REDIS_HOST, **redis_config) as rdb:
        rus = RedisUserSchema(rdb)
        yield rus
        
        await rus.close()
        
        
class JwtAuthHelper:
    def __init__(self,):
        self.__denylist = set()
        self.__redis_jti_key = 'auth:jwt:denyset'
        
    def extract_token(self, access_token:str)->TokenPayload|None:
        try:
            payload = jwt.decode(access_token, settings.AUTH_JWT_SECRET_KEY, algorithms=[settings.AUTH_ALGORITHM])

            return TokenPayload(**payload)
        except(jwt.JWTError, ValidationError):
            return None

    def extract_refresh_token(self, refresh_token:str)->TokenPayload|None:
        try:
            payload = jwt.decode(refresh_token, settings.AUTH_JWT_REFRESH_SECRET_KEY, algorithms=[settings.AUTH_ALGORITHM])

            return TokenPayload(**payload)
        except(jwt.JWTError, ValidationError):
            return None
        
    def generate_passport(self, email:str)->TokenSchema:
        return TokenSchema(access_token=JwtAuthHelper._create_access_token(email), refresh_token=JwtAuthHelper._create_refresh_token(email),)
    
    async def validate_access_active(self, access_token:str, redis_client:AioRedisCache=None):
        token_data = self.extract_token(access_token)
        # jti是否在黑名单中@redis
        if redis_client is not None:
            is_exists:bool = await redis_client.exist_set_member(self.__redis_jti_key, token_data.jti)
            if is_exists:
                raise HTTPException(status_code = status.HTTP_401_UNAUTHORIZED, detail="Token refresh", headers={"WWW-Authenticate": "Bearer"},)
                
        if token_data.jti in self.__denylist:
            raise HTTPException(status_code = status.HTTP_401_UNAUTHORIZED, detail="Token refresh", headers={"WWW-Authenticate": "Bearer"},)
            
        if datetime.fromtimestamp(token_data.exp) < datetime.now():
            raise HTTPException(status_code = status.HTTP_401_UNAUTHORIZED, detail="Token expired", headers={"WWW-Authenticate": "Bearer"},)
        
        return token_data
        
    def validate_refresh_active(self, refresh_token:str):
        token_data = self.extract_refresh_token(refresh_token)
        if datetime.fromtimestamp(token_data.exp) < datetime.now():
            raise HTTPException(status_code = status.HTTP_401_UNAUTHORIZED, detail="Token expired", headers={"WWW-Authenticate": "Bearer"},)
        return token_data
        
        
    async def refresh_passport(self, access_token:str, email:str, redis_client:AioRedisCache=None)->RefreshToken:
        token_data = self.extract_token(access_token)
        err_mesg = 'access token validate fail'
        if token_data is not None:
            if token_data.jti not in self.__denylist:
                if redis_client is not None:
                    await redis_client.add_set_key(self.__redis_jti_key, token_data.jti)
                    
                self.__denylist.add(token_data.jti)
                return RefreshToken(access_token=JwtAuthHelper._create_access_token(email), code=200)
            else:
                err_mesg = 'disable Repeat refresh'
            
        return RefreshToken(code=500, message=err_mesg)
            
    @staticmethod
    def _get_jwt_identifier() -> str:
        return str(uuid.uuid4())

    @staticmethod
    def _create_access_token(subject: Union[str, Any], expires_delta: int = None) -> str:
        if expires_delta is not None:
            expires_delta = datetime.utcnow() + expires_delta
        else:
            expires_delta = datetime.utcnow() + timedelta(minutes=settings.AUTH_ACCESS_TOKEN_EXPIRE_MINUTES)
    
        to_encode = {"exp": expires_delta, "sub": str(subject), "jti": JwtAuthHelper._get_jwt_identifier()}
        encoded_jwt = jwt.encode(to_encode, settings.AUTH_JWT_SECRET_KEY, settings.AUTH_ALGORITHM)
        return encoded_jwt

    @staticmethod
    def _create_refresh_token(subject: Union[str, Any], expires_delta: int = None) -> str:
        if expires_delta is not None:
            expires_delta = datetime.utcnow() + expires_delta
        else:
            expires_delta = datetime.utcnow() + timedelta(minutes=settings.AUTH_REFRESH_TOKEN_EXPIRE_MINUTES)
    
        to_encode = {"exp": expires_delta, "sub": str(subject), "jti": JwtAuthHelper._get_jwt_identifier()}
        encoded_jwt = jwt.encode(to_encode, settings.AUTH_JWT_REFRESH_SECRET_KEY, settings.AUTH_ALGORITHM)
        return encoded_jwt
    
jwt_auth_handler = JwtAuthHelper()
reuseable_oauth = ComposeOAuth2PasswordBearer(
    tokenUrl="/auth/login",
    scheme_name="JWT"
)

async def get_current_user(token: str = Depends(reuseable_oauth), db:RedisUserSchema=Depends(deps_user_db)) -> SystemUser:
    try:
        token_data = await jwt_auth_handler.validate_access_active(token, db.redis_client())
    except(jwt.JWTError, ValidationError, AttributeError):
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Could not validate credentials", headers={"WWW-Authenticate": "Bearer"},)
        
    userDict = await db.get_by_email(token_data.sub, defValue=None)
    
    if userDict is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Could not find user",)

    return SystemUser(**userDict)


    
async def validate_refresh_token(token: str = Depends(reuseable_oauth), db:RedisUserSchema=Depends(deps_user_db)) -> Tuple[SystemUser, RedisUserSchema]:
    try:
        token_data = jwt_auth_handler.validate_refresh_active(token)
    except(jwt.JWTError, ValidationError, AttributeError):
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Could not validate credentials", headers={"WWW-Authenticate": "Bearer"},)
        
    userDict = await db.get_by_email(token_data.sub, defValue=None)
    
    if userDict is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Could not find user",)
    return SystemUser(**userDict), db.redis_client()